import cv2
import os
import torch
import torch
import urllib.request
import cv2
import numpy as np
from torchvision.transforms import Compose, Resize, ToTensor, Normalize

def generate_depth_image_1(img):
	'''生成深度图像'''
	# 加载深度神经网络
	model = torch.hub.load('intel-isl/MiDaS', 'MiDaS')
	model.eval()

	# 图像预处理
	transform = Compose([
		Resize((384, 384)),
		ToTensor(),
		Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
	])
	img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
	img = transform(img).unsqueeze(0)

def generate_depth_image(img):
	img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)

	# 加载模型
	midas = torch.hub.load("intel-isl/MiDaS", "MiDaS_small")
	midas.eval()

	# 预处理器
	transform = torch.hub.load("intel-isl/MiDaS", "transforms").small_transform

	input_batch = transform(img).unsqueeze(0)

	# 预测深度
	with torch.no_grad():
			prediction = midas(input_batch)
			depth_map = prediction.squeeze().cpu().numpy()

	# 显示深度图
	depth_map = cv2.normalize(depth_map, None, 0, 255, cv2.NORM_MINMAX)
	depth_map = depth_map.astype(np.uint8)
	return depth_map

def save_frames_and_generate_depth(video_path):
	# 确保视频文件存在
	if not os.path.exists(video_path):
			print(f"视频文件 {video_path} 不存在")
			return

	# 获取视频文件的目录
	video_dir = os.path.dirname(video_path)
	images_dir = os.path.join(video_dir, 'images')
	depth_dir = os.path.join(video_dir, 'depth')

	# 创建子目录
	if not os.path.exists(images_dir):
			os.makedirs(images_dir)
	if not os.path.exists(depth_dir):
			os.makedirs(depth_dir)


	# 打开视频文件
	cap = cv2.VideoCapture(video_path)

	if not cap.isOpened():
			print(f"无法打开视频文件 {video_path}")
			return

	frame_count = 0

	while True:
			ret, frame = cap.read()
			if not ret:
					break

			# 保存帧为图片
			frame_filename = os.path.join(images_dir, f'frame-{frame_count:04d}.jpg')
			cv2.imwrite(frame_filename, frame)

			# 生成深度图像
			depth_image = generate_depth_image(frame)

			# 保存深度图像
			depth_filename = os.path.join(depth_dir, f'depth-{frame_count:04d}.jpg')
			cv2.imwrite(depth_filename, depth_image)

			frame_count += 1

	cap.release()
	print(f"已完成：共保存了 {frame_count} 帧图片和对应的深度图像")

# 示例调用
# save_frames_and_generate_depth('path/to/your/video.mp4')
